import requests
import json
import time

print("=== Dify OpenAI兼容API连接测试工具 ===\n")

# 测试不同的基础URL组合
url_combinations = [
    # 标准OpenAI格式
    {"base": "http://localhost:8006", "endpoint": "/v1/chat/completions", "name": "localhost不带结尾斜杠"},
    {"base": "http://127.0.0.1:8006", "endpoint": "/v1/chat/completions", "name": "127.0.0.1不带结尾斜杠"},
    
    # 带结尾斜杠的基础URL
    {"base": "http://localhost:8006/", "endpoint": "v1/chat/completions", "name": "localhost带结尾斜杠"},
    {"base": "http://127.0.0.1:8006/", "endpoint": "v1/chat/completions", "name": "127.0.0.1带结尾斜杠"},
    
    # 将v1作为基础URL的一部分
    {"base": "http://localhost:8006/v1", "endpoint": "/chat/completions", "name": "localhost带v1路径"},
    {"base": "http://127.0.0.1:8006/v1", "endpoint": "/chat/completions", "name": "127.0.0.1带v1路径"},
    
    # 测试Windows可能的特殊格式
    {"base": "http://localhost:8006", "endpoint": "/v1%20/chat/completions", "name": "测试URL空格问题"},
    
    # 测试不带/v1的直接路径
    {"base": "http://localhost:8006", "endpoint": "/chat/completions", "name": "直接chat路径(localhost)"},
    {"base": "http://127.0.0.1:8006", "endpoint": "/chat/completions", "name": "直接chat路径(127.0.0.1)"},
]

# 测试数据 - 符合OpenAI标准格式
test_data = {
    "model": "gemma-3-27b-it",
    "messages": [{"role": "user", "content": "你好，这是Dify连接测试"}],
    "temperature": 0.7,
    "max_tokens": 500
}

# 保存成功的结果
success_results = []

# 测试所有组合
print("\n正在测试聊天接口...")
print("-" * 60)

for combo in url_combinations:
    base = combo["base"]
    endpoint = combo["endpoint"]
    name = combo["name"]
    
    # 构建完整URL
    if endpoint.startswith("/"):
        full_url = f"{base}{endpoint}"
    else:
        full_url = f"{base}/{endpoint}"
    
    # 缩短用于显示的URL
    display_url = full_url if len(full_url) < 60 else f"{full_url[:57]}..."
    
    print(f"测试 [{name}]: {display_url}")
    
    try:
        # 添加API Key - Dify会传递Bearer token
        headers = {
            "Authorization": "Bearer dummy-key",
            "Content-Type": "application/json"
        }
        
        # 测量响应时间
        start_time = time.time()
        response = requests.post(full_url, json=test_data, headers=headers, timeout=10)
        elapsed = time.time() - start_time
        
        status = response.status_code
        print(f"  状态码: {status} (响应时间: {elapsed:.2f}秒)")
        
        if status == 200:
            # 提取响应内容
            try:
                result = response.json()
                content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
                print(f"  ✓ 成功! 响应内容: {content[:50]}...\n")
                
                # 保存成功结果，便于后续对比
                dify_base_url = base
                if "/v1" in endpoint and "/v1" not in base:
                    dify_base_url = f"{base}/v1"
                elif "/v1" not in endpoint and "/v1" not in base:
                    dify_base_url = f"{base}/v1"
                    
                success_results.append({
                    "name": name,
                    "tested_url": full_url,
                    "dify_base_url": dify_base_url.rstrip("/"),
                    "response_time": elapsed
                })
            except json.JSONDecodeError:
                print(f"  ⚠ 成功但JSON解析失败: {response.text[:100]}...\n")
        else:
            print(f"  ✗ 失败: {response.text[:100]}...\n")
    except Exception as e:
        print(f"  ✗ 错误: {str(e)}\n")

# 输出测试结果摘要
print("\n=== 测试结果摘要 ===")
if success_results:
    # 按响应时间排序，推荐最快的配置
    success_results.sort(key=lambda x: x["response_time"])
    best_result = success_results[0]
    
    print(f"找到 {len(success_results)} 个可用配置，推荐最快的配置:")
    print(f"配置名称: {best_result['name']}")
    print(f"实际测试URL: {best_result['tested_url']}")
    print(f"响应时间: {best_result['response_time']:.2f}秒")
    
    print("\n=== Dify配置指南 ===")
    print("1. 在Dify中选择'OpenAI API Compatible'")
    print("2. 按照下面的值填写配置:")
    print("   模型名称: gemma-3-27b-it")
    print("   API Key: dummy-key")
    print(f"   API Base URL: {best_result['dify_base_url']}")
    print("\n3. 如果仍然无法连接，尝试这些备选配置:")
    
    for i, result in enumerate(success_results[1:4], 1):  # 最多列出3个备选
        if i <= len(success_results) - 1:
            print(f"   备选{i}: {result['dify_base_url']} ({result['name']})")
    
    print("\n4. 重要提示:")
    print("   - 确保URL中没有多余空格")
    print("   - 如果Dify在Docker中运行，使用服务器IP替代localhost")
    print("   - 确保端口8006已在防火墙中开放")
else:
    print("没有找到可用的配置。请检查:")
    print("1. API服务是否正在运行 (python main.py)")
    print("2. 防火墙是否允许端口8006")
    print("3. 如果Dify在Docker中运行，需要使用主机IP而非localhost")